{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env SERPER_API_KEY=YOUR_SERPER_API_KEY\n",
    "%env OPENAI_API_KEY=YOUR_OPENAI_API_KEY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up\n",
    "\n",
    "The aim of this notebook is to showcase how one can use Langchain Agents using `Avatar` Module and optimize the actor using `AvatarOptimizer` optimizer for each of the toolset for the datasets. We'll be testing our module over 3 datasets:\n",
    "\n",
    "* ArxivQA\n",
    "* SearchAQ\n",
    "\n",
    "Before loading our datasets and going to the execution part, we'll need to configure the `lm` in `dspy.settings`. For the purpose of this notebook we'll be using `gpt-4o-mini`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import dspy\n",
    "\n",
    "dspy.settings.configure(\n",
    "    lm=dspy.OpenAI(\n",
    "        model=\"gpt-4o-mini\",\n",
    "        api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
    "        max_tokens=4000,\n",
    "        temperature=0,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining Signature\n",
    "\n",
    "Over all the three datasets the nature of problem is essentially a QA type so we'll create similar signatures `SearchQASignature` and `ArxivQASignature`. The only difference between them is `ArxivQASignature` takes `paper_id` as input too. This is mainly for Arxiv API tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SearchQASignature(dspy.Signature):\n",
    "    \"\"\"You will be given a question. Your task is to answer the question.\"\"\"\n",
    "    \n",
    "    question: str = dspy.InputField(\n",
    "        prefix=\"Question:\",\n",
    "        desc=\"question to ask\",\n",
    "        format=lambda x: x.strip(),\n",
    "    )\n",
    "    answer: str = dspy.OutputField(\n",
    "        prefix=\"Answer:\",\n",
    "        desc=\"answer to the question\",\n",
    "    )\n",
    "\n",
    "class ArxivQASignature(dspy.Signature):\n",
    "    \"\"\"You will be given a question and an Arxiv Paper ID. Your task is to answer the question.\"\"\"\n",
    "    \n",
    "    question: str = dspy.InputField(\n",
    "        prefix=\"Question:\",\n",
    "        desc=\"question to ask\",\n",
    "        format=lambda x: x.strip(),\n",
    "    )\n",
    "    paper_id: str = dspy.InputField(\n",
    "        prefix=\"Paper ID:\",\n",
    "        desc=\"Arxiv Paper ID\",\n",
    "    )\n",
    "    answer: str = dspy.OutputField(\n",
    "        prefix=\"Answer:\",\n",
    "        desc=\"answer to the question\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Datasets\n",
    "\n",
    "We'll be loading three datasets to evaluate our model on them. We'll be using `searchqa` and `arxiv_qa` datasets for the purpose of this notebook. We can use DSPy `DataLoader` to load these datasets from HuggingFace to DSPy friendly format of list of `Example`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import sample\n",
    "from dspy.datasets import DataLoader\n",
    "\n",
    "dl = DataLoader()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "searchqa = dl.from_huggingface(\n",
    "    \"lucadiliello/searchqa\",\n",
    "    split=\"train\",\n",
    "    input_keys=(\"question\",),\n",
    ")\n",
    "\n",
    "arxiv_qa = dl.from_huggingface(\n",
    "    \"taesiri/arxiv_qa\",\n",
    "    split=\"train\",\n",
    "    input_keys=(\"question\", \"paper_id\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to demonstration purposes we'll operate on a subset of training and testing dataset. We'll be using 200 examples for training set and 100 examples for testing set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "# Set a random seed for reproducibility\n",
    "random.seed(42)\n",
    "\n",
    "\n",
    "sqa_train = [\n",
    "    dspy.Example(question=example.question, answer=\",\".join(example.answers)).with_inputs(\"question\")\n",
    "    for example in sample(searchqa, 200)\n",
    "]\n",
    "sqa_test = [\n",
    "    dspy.Example(question=example.question, answer=\",\".join(example.answers)).with_inputs(\"question\")\n",
    "    for example in sample(searchqa, 100)\n",
    "]\n",
    "\n",
    "aqa_train = [\n",
    "    dspy.Example(question=example.question, paper_id=example.paper_id, answer=example.answer).with_inputs(\"question\", \"paper_id\")\n",
    "    for example in sample(arxiv_qa, 200)\n",
    "]\n",
    "aqa_test = [\n",
    "    dspy.Example(question=example.question, paper_id=example.paper_id, answer=example.answer).with_inputs(\"question\", \"paper_id\")\n",
    "    for example in sample(arxiv_qa, 100)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up Tools\n",
    "\n",
    "We'll setup `Avatar` modules for both signatures and all the `tools` can be used by each of the dataset i.e. `searchqa` and `arxiv_qa`. `Tool` is a pydantic model that Avatar expects the `tools` to be composed as more specifically it have 4 fields:\n",
    "\n",
    "* `name` : Name of the tool\n",
    "* `input_type` : Type of input the tool accepts\n",
    "* `output_type` : Type of output the tool returns\n",
    "* `tool` : The actual tool object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dspy.predict.avatar import Tool, Avatar\n",
    "from langchain.tools import WikipediaQueryRun\n",
    "from langchain_community.utilities import GoogleSerperAPIWrapper, ArxivAPIWrapper, WikipediaAPIWrapper\n",
    "\n",
    "tools = [\n",
    "    Tool(\n",
    "        tool=GoogleSerperAPIWrapper(),\n",
    "        name=\"WEB_SEARCH\",\n",
    "        desc=\"If you have a question, you can use this tool to search the web for the answer.\"\n",
    "    ),\n",
    "    Tool(\n",
    "        tool=ArxivAPIWrapper(),\n",
    "        name=\"ARXIV_SEARCH\",\n",
    "        desc=\"Pass the arxiv paper id to get the paper information.\",\n",
    "        input_type=\"Arxiv Paper ID\",\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have defined our `tools`, we can now create an `Avatar` object by passing the `tools` and `signature`. It takes 2 more optional parameters `verbose` and `max_iters`. `verbose` is used to display the logs and `max_iters` is used to control the number of iterations in multi step execution. \n",
    "\n",
    "An avatar agent stops the tool usage iteration once it reaches `max_iters` or when it prompts `Finish`. You can also create custom tools too, all you need to make sure is:\n",
    "\n",
    "* You pass is a class object.\n",
    "* Implements `__init__` and `run` method.\n",
    "* Must take 1 string a input and returns 1 string as output.\n",
    "\n",
    "If your tool doesn't return or takes input a string then you can make a custom wrapper to take care of that for now. In future we'll try to enable a diverse tool usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arxiv_agent = Avatar(\n",
    "    tools=tools,\n",
    "    signature=ArxivQASignature,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "search_agent = Avatar(\n",
    "    tools=tools,\n",
    "    signature=SearchQASignature,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "Open enden QA tasks are hard to evaluate on rigid metrics like exact match. So, we'll be using an improvised LLM as Judge for the evaluation of our model on test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Evaluator(dspy.Signature):\n",
    "    \"\"\"Please act as an impartial judge and evaluate the quality of the responses provided by multiple AI assistants to the user question displayed below. You should choose the assistant that offers a better user experience by interacting with the user more effectively and efficiently, and providing a correct final response to the user's question.\n",
    "    \n",
    "Rules:\n",
    "1. Avoid Position Biases: Ensure that the order in which the responses were presented does not influence your decision. Evaluate each response on its own merits.\n",
    "2. Length of Responses: Do not let the length of the responses affect your evaluation. Focus on the quality and relevance of the response. A good response is targeted and addresses the user's needs effectively, rather than simply being detailed.\n",
    "3. Objectivity: Be as objective as possible. Consider the user's perspective and overall experience with each assistant.\"\"\"\n",
    "    \n",
    "    question: str = dspy.InputField(\n",
    "        prefix=\"Question:\",\n",
    "        desc=\"question to ask\",\n",
    "    )\n",
    "    reference_answer: str = dspy.InputField(\n",
    "        prefix=\"Reference Answer:\",\n",
    "        desc=\"Answer to the question given by the model.\",\n",
    "    )\n",
    "    answer: str = dspy.InputField(\n",
    "        prefix=\"Answer:\",\n",
    "        desc=\"Answer to the question given by the model.\",\n",
    "    )\n",
    "    rationale: str = dspy.OutputField(\n",
    "        prefix=\"Rationale:\",\n",
    "        desc=\"Explanation of why the answer is correct or incorrect.\",\n",
    "    )\n",
    "    is_correct: bool = dspy.OutputField(\n",
    "        prefix=\"Correct:\",\n",
    "        desc=\"Whether the answer is correct.\",\n",
    "    )\n",
    "\n",
    "\n",
    "evaluator = dspy.TypedPredictor(Evaluator)\n",
    "\n",
    "\n",
    "def metric(example, prediction, trace=None):\n",
    "    return int(\n",
    "        evaluator(\n",
    "            question=example.question,\n",
    "            answer=prediction.answer,\n",
    "            reference_answer=example.answer\n",
    "        ).is_correct\n",
    "    ) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For evaluation we can't use `dspy.Evaluate`, reason being that `Avatar` changes it's signature per iteration by adding the actions and it's results to it as fields. So we can create our own hacky thread safe evaluator for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "\n",
    "def process_example(example, signature):\n",
    "    try:\n",
    "        avatar = Avatar(\n",
    "            signature,\n",
    "            tools=tools,\n",
    "            verbose=False,\n",
    "        )\n",
    "        prediction = avatar(**example.inputs().toDict())\n",
    "\n",
    "        return metric(example, prediction)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return 0\n",
    "\n",
    "\n",
    "def multi_thread_executor(test_set, signature, num_threads=60):\n",
    "    total_score = 0\n",
    "    total_examples = len(test_set)\n",
    "\n",
    "    with ThreadPoolExecutor(max_workers=num_threads) as executor:\n",
    "        futures = [executor.submit(process_example, example, signature) for example in test_set]\n",
    "\n",
    "        for future in tqdm.tqdm(futures, total=total_examples, desc=\"Processing examples\"):\n",
    "            total_score += future.result()\n",
    "\n",
    "    avg_metric = total_score / total_examples\n",
    "    return avg_metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sqa_score = multi_thread_executor(sqa_test, SearchQASignature)\n",
    "print(f\"Average Score on SearchQA: {sqa_score:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aqa_score = multi_thread_executor(aqa_test, ArxivQASignature)\n",
    "print(f\"Average Score on ArxivQA: {aqa_score:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization\n",
    "\n",
    "For the optimization of the `Actor` we'll be using `AvatarOptimizer`. It's a DSPy implementation of the [Avatar](https://github.com/zou-group/avatar/) method that optimizes the `Actor` for the given `tools` using a comparator module that optimizes Actor instruction. Note, that Actor is the Module that directs tool execution and flow, it's not the signature that we are passing. It doesn't optimize the instruction of the signature we pass. It takes the following parameters:\n",
    "\n",
    "* `metric`: Metric that we'll be optimizing for\n",
    "* `max_iters`: Maximum number of iterations for the optimizer\n",
    "* `lower_bound`: Lower bound for the metric to classify example as negative\n",
    "* `upper_bound`: Upper bound for the metric to classify example as positive\n",
    "* `max_positive_inputs`: Maximum number of positive inputs sampled for comparator\n",
    "* `max_negative_inputs`: Maximum number of negative inputs sampled for comparator\n",
    "* `optimize_for`: Whether we want to maximize the metric or minimize it during optimization\n",
    "\n",
    "Once the optimizer is done we can get the optimized actor and use it for the evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dspy.teleprompt import AvatarOptimizer\n",
    "\n",
    "teleprompter = AvatarOptimizer(\n",
    "    metric=metric,\n",
    "    max_iters=1,\n",
    "    max_negative_inputs=10,\n",
    "    max_positive_inputs=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_arxiv_agent = teleprompter.compile(\n",
    "    student=arxiv_agent,\n",
    "    trainset=aqa_train\n",
    ")\n",
    "\n",
    "optimized_search_agent = teleprompter.compile(\n",
    "    student=search_agent,\n",
    "    trainset=sqa_train\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can evaluate our actor module, for this we've provided an implementation of thread safe evaluator that we above as part of class method of `AvatarOptimizer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "teleprompter.thread_safe_evaluator(aqa_test, optimized_arxiv_agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "teleprompter.thread_safe_evaluator(sqa_test, optimized_search_agent)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dspy-74wouE_3-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
